from numbers import Number
from typing_extensions import Self
from unittest import result
import cv2
import numpy as np
from paddleocr import PaddleOCR


class AdaptiveInvoiceOCR:
    def __init__(self, use_gpu=False):
        self.ocr = PaddleOCR(use_angle_cls=True, lang="ch", use_gpu=use_gpu)

    def standardize_image(self, image_path, target_width=1206):
        """步骤1: 按宽度1200标准化图像，保持宽高比"""
        img = cv2.imread(image_path)
        if img is None:
            raise FileNotFoundError(f"无法读取图像: {image_path}")

        h, w = img.shape[:2]
        scale = target_width / w
        new_h = int(h * scale)
        img = cv2.resize(img, (target_width, new_h))
        return img, scale

    # def get_text_blocks(self, img):
    #     """获取所有文本块 (x, y, w, h, text)"""
    #     result = self.ocr.ocr(img, cls=True)
    #     blocks = []
    #     print(f'get_text_blocks:{result}')
    #     for line in result:
    #         for word_info in line:
    #             box = word_info[0]  # [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
    #             text = word_info[1][0]
    #             # 计算外接矩形
    #             x_coords = [p[0] for p in box]
    #             y_coords = [p[1] for p in box]
    #             x, y, w, h = cv2.boundingRect(np.array(box))
    #             blocks.append({
    #                 'x': x,
    #                 'y': y,
    #                 'w': w,
    #                 'h': h,
    #                 'text': text,
    #                 'box': box
    #             })
    #     # 按 y 排序（从上到下）
    #     blocks.sort(key=lambda b: b['y'])
    #     return blocks
    def get_text_blocks(self, img, minY, maxY):
        """获取所有文本块 (x, y, w, h, text)"""
        crop = img[minY:maxY, 0:1206]
        result = self.ocr.ocr(crop, cls=True)

        # print(f"get_text_blocks:{result}")
        # row = []
        # current = []
        blocks = []
        for line in result:
            for word_info in line:
                box = word_info[0]  # [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
                text = word_info[1][0]

                # 直接从 box 中计算最小外接矩形
                xs = [point[0] for point in box]
                ys = [point[1] for point in box]
                x_min, x_max = min(xs), max(xs)
                y_min, y_max = min(ys), max(ys)

                w = x_max - x_min
                h = y_max - y_min

                blocks.append(
                    {"x": x_min, "y": y_min, "w": w, "h": h, "text": text, "box": box}
                )
                # print(f"Text:{text}")

        # 按 y 排序（从上到下）
        blocks.sort(key=lambda b: b["y"])
        # print(f"cs:{cs}")
        return result[0]

    # def get_text_blocks(self, img):
    #     """获取所有文本块 (x, y, w, h, text)"""
    #     result = self.ocr.ocr(img, cls=True)
    #     blocks = []
    #     for line in result:
    #         for word_info in line:
    #             box = word_info[0]  # [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
    #             text = word_info[1][0]
    #             # 确保 box 是一个有效的点集
    #             if len(box) >= 4:  # 至少需要四个点来定义一个矩形
    #                 pts = np.array(box, dtype=np.int32)
    #                 rect = cv2.boundingRect(pts)
    #                 x, y, w, h = rect
    #                 blocks.append({
    #                     'x': x,
    #                     'y': y,
    #                     'w': w,
    #                     'h': h,
    #                     'text': text,
    #                     'box': box
    #                 })
    #     # 按 y 排序（从上到下）
    #     blocks.sort(key=lambda b: b['y'])
    #     return blocks

    def detect_header_footer(self, blocks, minY, maxY):
        """步骤3&4: 自适应划分头部、表格、尾部"""
        # header_bottom = int(img_h * header_ratio)  # 前35%为头部
        # footer_top = int(img_h * (1 - footer_ratio))  # 后20%为尾部

        # header_blocks = [b for b in blocks if b["y"] < header_bottom]
        # footer_blocks = [b for b in blocks if b["y"] > footer_top]
        table_blocks = [b for b in blocks if b["y"] >= minY and b["y"] <= maxY]

        return table_blocks

    def extract_table_structure(self, table_blocks, min_row_height=20):
        """步骤5: 从 table_blocks 中提取表格二维结构"""
        print(f"table_blocks:{table_blocks}")
        if not table_blocks:
            return []

        # 1. 按 y 分组为行
        table_blocks.sort(key=lambda b: (b["y"], b["x"]))
        rows = []
        current_row = []
        last_y = table_blocks[0]["y"]

        for block in table_blocks:
            if abs(block["y"] - last_y) > min_row_height:
                if current_row:
                    rows.append(current_row)
                current_row = [block]
                last_y = block["y"]
            else:
                current_row.append(block)
        if current_row:
            rows.append(current_row)

        # 2. 对每一行按 x 排序
        for row in rows:
            row.sort(key=lambda b: b["x"])

        # 3. 提取文本
        table_data = []
        for row in rows:
            print(f"row:{row}")
            row_texts = [b["text"] for b in row]
            table_data.append(row_texts)

        return table_data

    def extract_from_image(self, image_path):
        # 步骤1: 标准化图像
        img, scale = self.standardize_image(image_path, target_width=1200)
        img_h, img_w = img.shape[:2]

        # 步骤2: OCR 获取所有文本块
        tableMinY, tableMaxY = 300, img_h - 290
        all_blocks = self.get_text_blocks(img, tableMinY, tableMaxY)
        print(f"所有的文本块：all_blocks:{all_blocks}")

        # 步骤3&4: 划分头部、表格、尾部
        # table_blocks = self.detect_header_footer(all_blocks, tableMinY, tableMaxY)

        # 步骤5: 提取表格结构
        # table_data = self.extract_table_structure(all_blocks)
        table_data = self.ocr_result_to_2d_table(all_blocks, y_threshold=15)
        print(f"tableData:{table_data}")

        header_result_json = self.header_result(img)
        print(f"header_result_json:{header_result_json}")

        foot_result_json = self.foot_result(img_h, img)
        print(f"foot_result_json:{foot_result_json}")

        # 可选：提取头部和尾部文本（可用于发票类型判断）
        # header_text = " ".join([b["text"] for b in header_blocks])
        # footer_text = " ".join([b["text"] for b in footer_blocks])

        # 返回结构化结果
        # result = {
        #     "header": header_text,
        #     "items": table_data,  # 二维数组
        #     "footer": footer_text,
        #     "image_height": img_h,
        #     "scale": scale,
        # }

        result = {**header_result_json, **foot_result_json, "items": table_data}

        return result

    def foot_result(self, img_h, img):
        foot_json = {
            "total": {
                "amount_total": {
                    "x": 757,
                    "y": img_h - 283,
                    "w": 908,
                    "h": img_h - 253,
                },
                "tax_total": {"x": 1059, "y": img_h - 283, "w": 1180, "h": img_h - 253},
                "total_in_words": {
                    "x": 300,
                    "y": img_h - 246,
                    "w": 836,
                    "h": img_h - 206,
                },
                "total_in_digits": {
                    "x": 893,
                    "y": img_h - 263,
                    "w": 1179,
                    "h": img_h - 206,
                },
            },
            "remark": {"x": 65, "y": img_h - 210, "w": 1177, "h": img_h - 87},
            "open": {"x": 171, "y": img_h - 72, "w": 943, "h": img_h - 30},
        }
        result = self._extract_recursive(foot_json, img)
        # print(f"FootResult:{result}")
        return result

    def header_result(self, img):
        header_json = {
            "invoice_type": {"x": 261, "y": 28, "w": 868, "h": 92},
            "invoice_number": {"x": 969, "y": 57, "w": 1204, "h": 89},
            "issue_date": {"x": 970, "y": 98, "w": 1206, "h": 129},
            "buyer": {
                "name": {"x": 65, "y": 176, "w": 600, "h": 235},
                "code": {"x": 335, "y": 234, "w": 600, "h": 299},
            },
            "seller": {
                "name": {"x": 642, "y": 176, "w": 1178, "h": 235},
                "code": {"x": 911, "y": 234, "w": 1178, "h": 299},
            },
        }
        result = self._extract_recursive(header_json, img)
        # print(f"HeaderResult:{result}")
        return result

    def _extract_recursive(self, template, img):
        """
        递归提取模板中的所有文本
        """
        result = {}
        for key, value in template.items():
            if isinstance(value, dict):
                if "x" in value and "y" in value and "w" in value and "h" in value:
                    # 是叶子节点（有坐标)
                    x, y, w, h = value["x"], value["y"], value["w"], value["h"]
                    text = self.recognize_region(img, x, y, w, h)
                    # print(f"text:{text}")

                    if "名称：" in text:
                        text = text.replace("名称：", "")
                    if "统一社会信用代码/纳税人识别号：" in text:
                        text = text.replace("统一社会信用代码/纳税人识别号：", "")

                    result[key] = text
                else:
                    # 是中间节点（嵌套结构）
                    zj = self._extract_recursive(value, img)
                    # print(f"Key:{key},中间节点:{zj}")
                    result[key] = zj
            else:
                # 兜底：非 dict 直接跳过（理论上不会出现）
                result[key] = str(value)
        return result

    def recognize_region(self, img, x, y, w, h):
        """
        裁剪区域并 OCR 识别，支持多行文本合并
        """
        # 确保坐标和尺寸为整数
        x, y, w, h = int(x), int(y), int(w), int(h)
        # 边界检查
        # h_img, w_img = img.shape[:2]
        # x = max(0, x)
        # y = max(0, y)
        # w = min(w, w_img - x)
        # h = min(h, h_img - y)
        # if w <= 0 or h <= 0:
        #     return ""

        crop = img[y:h, x:w]
        # print(f"Crop:{crop}...x:{x},w:{w},y:{y},h:{h}")
        result = self.ocr.ocr(crop, cls=True)
        texts = []
        for line in result:
            for word_info in line:
                word = word_info[1][0]  # 提取文本
                texts.append(word.strip())
        return " ".join(texts)

    # def ocr_result_to_2d_table(self, ocr_result, y_threshold=10):
    #     """
    #     将 PaddleOCR 的表格区域识别结果转换为二维表格
    #     :param ocr_result: PaddleOCR 输出的 result[0]（即单页结果）
    #     :param y_threshold: 判断是否为同一行的 y 坐标容差（像素）
    #     :return: 二维列表 [[...], [...], ...]
    #     """
    #     if not ocr_result:
    #         return []

    #     # 1. 提取每个文本块的信息：文本、中心 y、中心 x
    #     blocks = []
    #     for item in ocr_result:
    #         box = item[0]  # 四个点 [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
    #         text = item[1][0]  # 识别文本

    #         # 计算外接矩形中心 y 和 x
    #         xs = [p[0] for p in box]
    #         ys = [p[1] for p in box]
    #         center_y = (min(ys) + max(ys)) / 2
    #         center_x = (min(xs) + max(xs)) / 2

    #         blocks.append({"text": text, "center_y": center_y, "center_x": center_x})

    #     # 2. 按 center_y 排序
    #     blocks.sort(key=lambda b: b["center_y"])

    #     # 3. 按 y 聚类分组为行
    #     rows = []
    #     current_row = []
    #     last_y = blocks[0]["center_y"]

    #     for block in blocks:
    #         if abs(block["center_y"] - last_y) < y_threshold:
    #             current_row.append(block)
    #         else:
    #             if current_row:
    #                 rows.append(current_row)
    #             current_row = [block]
    #             last_y = block["center_y"]
    #     if current_row:
    #         rows.append(current_row)

    #     # 4. 每行按 x 排序，并提取文本
    #     table_data = []
    #     for row in rows:
    #         row.sort(key=lambda b: b["center_x"])  # 从左到右
    #         row_texts = [b["text"] for b in row]
    #         if row_texts is not Number:
    #             row_texts = self.merge_short_texts_in_row(row_texts, 2)
    #         table_data.append(row_texts)

    #     return table_data

    # def ocr_result_to_2d_table(self, ocr_result, y_threshold=15):
    #     """
    #     高级表格重建：支持列对齐、空值填充、跨行文本合并
    #     """
    #     if not ocr_result:
    #         return []

    #     # === 1. 提取所有文本块信息 ===
    #     blocks = []

    #     for item in ocr_result:
    #         try:
    #             box = item[0]
    #             text = item[1][0].strip()
    #         except:
    #             continue

    #         if not (isinstance(box, (list, tuple)) and len(box) > 0):
    #             continue
    #         if not isinstance(box[0], (list, tuple)) or len(box[0]) < 2:
    #             continue

    #         ys = [float(p[1]) for p in box if len(p) >= 2]
    #         xs = [float(p[0]) for p in box if len(p) >= 2]
    #         if not xs or not ys:
    #             continue

    #         center_y = (min(ys) + max(ys)) / 2
    #         center_x = (min(xs) + max(xs)) / 2
    #         left_x = min(xs)
    #         right_x = max(xs)
    #         height = max(ys) - min(ys)

    #         blocks.append(
    #             {
    #                 "text": text,
    #                 "center_y": center_y,
    #                 "center_x": center_x,
    #                 "left_x": left_x,
    #                 "right_x": right_x,
    #                 "height": height,
    #             }
    #         )

    #     if not blocks:
    #         return []

    #     # === 2. 按 y 分组为行，并合并碎片文本 ===
    #     blocks.sort(key=lambda b: b["center_y"])
    #     rows = []
    #     current_row = []
    #     last_y = blocks[0]["center_y"]

    #     for block in blocks:
    #         if abs(block["center_y"] - last_y) < y_threshold:
    #             current_row.append(block)
    #         else:
    #             if current_row:
    #                 # ✅ 在这里合并碎片文本
    #                 merged_row = self.merge_nearby_blocks_in_line(current_row, x_gap=25)
    #                 rows.append(merged_row)
    #             current_row = [block]
    #             last_y = block["center_y"]

    #     if current_row:
    #         merged_row = self.merge_nearby_blocks_in_line(current_row, x_gap=25)
    #         rows.append(merged_row)

    #     if not rows:
    #         return []
    #     # === 3. 提取第一行作为“列头”，确定列边界 ===
    #     header_row = rows[0]
    #     header_row.sort(key=lambda b: b["left_x"])  # 从左到右

    #     # 生成列边界：每个列的 x 范围 [start, end]
    #     col_borders = []
    #     for i, block in enumerate(header_row):
    #         if i == 0:
    #             start = 0
    #         else:
    #             # 列开始 = 上一列右侧 和 当前列左侧 的中间
    #             prev_right = header_row[i - 1]["right_x"]
    #             curr_left = block["left_x"]
    #             start = (prev_right + curr_left) / 2
    #         if i == len(header_row) - 1:
    #             end = float("inf")
    #         else:
    #             next_left = header_row[i + 1]["left_x"]
    #             end = (block["right_x"] + next_left) / 2
    #         col_borders.append((start, end))

    #     num_cols = len(col_borders)

    #     # === 4. 逐行重建表格，支持跨行合并 ===
    #     table_data = []
    #     last_row_items = [""] * num_cols  # 上一行内容

    #     for i, row in enumerate(rows):
    #         row.sort(key=lambda b: b["left_x"])
    #         current_row = [""] * num_cols

    #         # ✅ 判断是否是新行：以 * 开头 或 第一列文本以 * 开头
    #         is_new_line = any(b["text"].startswith("*") for b in row)

    #         # ✅ 新增：如果第一列文本不以 * 开头，且不是第一行，则可能是延续
    #         first_col_text = ""
    #         for block in row:
    #             if block["left_x"] == min(b["left_x"] for b in row):
    #                 first_col_text = block["text"]
    #                 break

    #         is_continuation = i > 0 and not first_col_text.startswith("*")

    #         if is_continuation:
    #             # 将当前行所有文本追加到上一行对应列
    #             for block in row:
    #                 col_idx = -1
    #                 for c, (start, end) in enumerate(col_borders):
    #                     if start <= block["center_x"] <= end:
    #                         col_idx = c
    #                         break
    #                 if 0 <= col_idx < num_cols:
    #                     last_row_items[col_idx] += block["text"]
    #             # 不生成新行
    #             continue
    #         else:
    #             # 正常新行
    #             for block in row:
    #                 col_idx = -1
    #                 for c, (start, end) in enumerate(col_borders):
    #                     if start <= block["center_x"] <= end:
    #                         col_idx = c
    #                         break
    #                 if 0 <= col_idx < num_cols:
    #                     current_row[col_idx] = block["text"]
    #             table_data.append(current_row)
    #             last_row_items = current_row[:]
    #     return table_data

    def ocr_result_to_2d_table(
        self, ocr_result, y_threshold=15, x_gap=25, col_threshold=20
    ):
        """
        将 OCR 结果转换为结构化二维表格，支持：
        - 碎片文本合并（如 '数'+'量' → '数量'）
        - 坐标合并
        - 列对齐与空值填充
        - 跨行文本合并（非 * 开头的第一列内容合并到上一行）

        :param ocr_result: PaddleOCR 输出 result[0]
        :param y_threshold: 行间距判断阈值（像素）
        :param x_gap: 碎片文本合并的最大 x 间隙
        :param col_threshold: 列对齐容差（影响列边界划分）
        :return: 二维列表，每行代表一行表格数据
        """
        if not ocr_result:
            return []

        # === 1. 提取所有文本块 ===
        blocks = []
        for item in ocr_result:
            try:
                box = item[0]
                text = item[1][0].strip()
                if not text:
                    continue
            except Exception:
                continue

            # 提取坐标
            if not (
                isinstance(box, (list, tuple))
                and len(box) > 0
                and isinstance(box[0], (list, tuple))
            ):
                continue
            try:
                ys = [float(p[1]) for p in box if len(p) >= 2]
                xs = [float(p[0]) for p in box if len(p) >= 2]
                if not xs or not ys:
                    continue
            except Exception:
                continue

            center_y = (min(ys) + max(ys)) / 2
            center_x = (min(xs) + max(xs)) / 2
            left_x = min(xs)
            right_x = max(xs)
            height = max(ys) - min(ys)

            blocks.append(
                {
                    "text": text,
                    "center_y": center_y,
                    "center_x": center_x,
                    "left_x": left_x,
                    "right_x": right_x,
                    "height": height,
                }
            )

        if not blocks:
            return []

        # === 2. 按 y 分组为行，并合并碎片文本 ===
        def merge_nearby_blocks_in_line(line_blocks, gap=x_gap):
            if len(line_blocks) <= 1:
                return line_blocks
            sorted_blocks = sorted(line_blocks, key=lambda b: b["left_x"])
            merged = [sorted_blocks[0]]
            for curr in sorted_blocks[1:]:
                last = merged[-1]
                if curr["left_x"] <= last["right_x"] + gap:  # 可合并
                    merged[-1] = {
                        "text": last["text"] + curr["text"],
                        "left_x": last["left_x"],
                        "right_x": max(last["right_x"], curr["right_x"]),
                        "center_x": (last["left_x"] + curr["right_x"]) / 2,
                        "center_y": (last["center_y"] + curr["center_y"]) / 2,
                        "height": max(last["height"], curr["height"]),
                    }
                else:
                    merged.append(curr)
            return merged

        blocks.sort(key=lambda b: b["center_y"])
        rows = []  # 存储合并后的行
        current_row = []
        last_y = blocks[0]["center_y"]

        for block in blocks:
            if abs(block["center_y"] - last_y) < y_threshold:
                current_row.append(block)
            else:
                if current_row:
                    merged_line = merge_nearby_blocks_in_line(current_row)
                    rows.append(merged_line)
                current_row = [block]
                last_y = block["center_y"]

        if current_row:
            merged_line = merge_nearby_blocks_in_line(current_row)
            rows.append(merged_line)

        if not rows:
            return []

        # === 3. 使用第一行确定列边界 ===
        header_row = rows[0]
        header_row.sort(key=lambda b: b["left_x"])
        col_borders = []  # 每列的 [start, end] x 范围

        for i, block in enumerate(header_row):
            if i == 0:
                start = 0
            else:
                prev_end = header_row[i - 1]["right_x"]
                curr_start = block["left_x"]
                start = (prev_end + curr_start) / 2

            if i == len(header_row) - 1:
                end = float("inf")
            else:
                curr_end = block["right_x"]
                next_start = header_row[i + 1]["left_x"]
                end = (curr_end + next_start) / 2

            col_borders.append((start, end))

        num_cols = len(col_borders)

        # === 4. 逐行重建表格，智能判断是否为延续行（修正版）===
        table_data = []
        last_row_items = [""] * num_cols
        last_row_y = None
        last_row_height = None

        for row_idx, row in enumerate(rows):
            row_sorted = sorted(row, key=lambda b: b["left_x"])
            current_row = [""] * num_cols

            # === 智能判断是否为延续行 ===
            is_continuation = False

            if row_idx > 0 and last_row_y is not None:
                first_block = row_sorted[0]
                curr_y = first_block["center_y"]
                curr_height = first_block["height"]
                curr_text = first_block["text"].strip()

                y_dist = abs(curr_y - last_row_y)
                max_allow_gap = max(last_row_height, curr_height) * 2.5

                is_short_text = len(curr_text) <= 4
                starts_with_star = curr_text.startswith("*")

                prev_first_col_range = (
                    col_borders[0] if col_borders else (0, float("inf"))
                )
                x_aligned = (
                    prev_first_col_range[0]
                    <= first_block["center_x"]
                    <= prev_first_col_range[1] * 1.3
                )
                # 是否以 * 开头是判断新列的唯一标准
                if (
                    # y_dist < max_allow_gap
                    # and is_short_text
                    # and
                      not starts_with_star
                    # and x_aligned
                ):
                    is_continuation = True

            if is_continuation:
                # ✅ 重点：修改 table_data 的最后一行
                last_table_row = table_data[-1]
                for block in row_sorted:
                    col_idx = -1
                    for c, (start, end) in enumerate(col_borders):
                        if start <= block["center_x"] <= end:
                            col_idx = c
                            break
                    if 0 <= col_idx < num_cols:
                        last_table_row[col_idx] += block["text"]  # 追加文本
                # ✅ 同步 last_row_items
                last_row_items = last_table_row[:]
                continue
            else:
                # 正常新行
                for block in row_sorted:
                    col_idx = -1
                    for c, (start, end) in enumerate(col_borders):
                        if start <= block["center_x"] <= end:
                            col_idx = c
                            break
                    if 0 <= col_idx < num_cols:
                        current_row[col_idx] = block["text"]
                table_data.append(current_row)
                last_row_items = current_row[:]
                last_row_y = row_sorted[0]["center_y"]
                last_row_height = row_sorted[0]["height"]

        return table_data

    def merge_short_texts_in_row(self, row, min_len=2):
        """合并长度小于 min_len 的连续文本"""
        result = []
        temp = ""
        for text in row:
            if len(text) < min_len:
                temp += text
            else:
                if temp:
                    result.append(temp)
                    temp = ""
                result.append(text)
        if temp:
            result.append(temp)
        return result

    def merge_nearby_blocks_in_line(self, blocks, x_gap=30):
        """
        将同一行中距离很近的文本块合并（如 '数' '量' → '数量'）
        :param blocks: 同一行的文本块列表
        :param x_gap: 最大允许间隙（像素），超过则不合并
        :return: 合并后的新列表
        """
        if len(blocks) <= 1:
            return blocks

        # 按 x 从左到右排序
        blocks_sorted = sorted(blocks, key=lambda b: b["left_x"])
        merged = [blocks_sorted[0]]

        for curr in blocks_sorted[1:]:
            last = merged[-1]
            # 判断是否紧邻（当前左边界 < 上一个右边界 + 间隙）
            if curr["left_x"] <= last["right_x"] + x_gap:
                # 合并文本和坐标
                merged[-1] = {
                    "text": last["text"] + curr["text"],
                    "left_x": last["left_x"],
                    "right_x": max(last["right_x"], curr["right_x"]),
                    "center_x": (last["left_x"] + curr["right_x"]) / 2,
                    "center_y": (last["center_y"] + curr["center_y"]) / 2,
                    "height": max(last["height"], curr["height"]),
                }
            else:
                merged.append(curr)
        return merged
